/**
 * Java implementation of a Binary Search Tree of type T.
 * 
 * @author Jonathan E. Landrum <me@jonlandrum.com>
 * @since 2016-01-08
 */

package com.jonlandrum.selfbalancingtree;

import java.util.Iterator;
import java.util.LinkedList;

public class BinarySearchTree<T extends Comparable<T>> {
    /*
     * Instance variables
     */
    protected int numElements = 0;
    protected Node<T> root = new Node<T>();
    
    /*
     * Constructors
     */
    public BinarySearchTree() {}
    
    public BinarySearchTree(T d) {
        this.addElement(new Node<T>(d));
    }
    
    public BinarySearchTree(Node<T> n) {
        this.addElement(n);
    }
    
    /*
     * Getters
     */
    public Node<T> getRoot() {
        return root;
    }
    
    public int getNumElements() {
        return numElements;
    }
    
    /*
     * Booleans
     */
    public boolean hasRoot() {
        return !root.isEmpty();
    }
    
    public boolean hasElement(T d) {
        return this.hasElement(new Node<T>(d));
    }
    
    public boolean hasElement(Node<T> n) {
        return !this.findElement(n).isEmpty();
    }
    
    /*
     * Operations
     */
    public Node<T> addElement(T d) {
        return this.addElement(new Node<T>(d));
    }
    
    public Node<T> addElement(Node<T> n) {
        Node<T> result = new Node<T>();
        if (n.isEmpty()) {
            return result;
        } else if (!this.hasRoot()) {
            root = n;
            result = root;
        } else {
            Node<T> current = root;
            while (true) {
                if (current.compareTo(n) < 0) {
                    if (current.hasRightChild()) {
                        current = current.getRightChild();
                    } else {
                        current.setRightChild(n);
                        n.setParent(current);
                        result = n;
                        break;
                    }
                } else {
                    if (current.hasLeftChild()) {
                        current = current.getLeftChild();
                    } else {
                        current.setLeftChild(n);
                        n.setParent(current);
                        result = n;
                        break;
                    }
                }
            }
        }
        ++numElements;
        return result;
    }
    
    public Node<T> findElement(T d) {
        return this.findElement(new Node<T>(d));
    }
    
    public Node<T> findElement(Node<T> n) {
        Node<T> result = root;
        if (n.isEmpty()) {
            result = new Node<T>();
        } else if (!result.isEmpty()) {
            while (result.compareTo(n) != 0) {
                if (result.compareTo(n) < 0) {
                    if (result.hasRightChild()) {
                        result = result.getRightChild();
                    } else {
                        result = new Node<T>();
                        break;
                    }
                } else {
                    if (result.hasLeftChild()) {
                        result = result.getLeftChild();
                    } else {
                        result = new Node<T>();
                        break;
                    }
                }
            }
        }
        return result;
    }
    
    public Node<T> removeElement(T d) {
        return this.removeElement(new Node<T>(d));
    }
    
    public Node<T> removeElement(Node<T> n) {
        Node<T> result = this.findElement(n);
        if (n.isEmpty() || result.isEmpty()) { return result; }
        if (!result.hasLeftChild() && !result.hasRightChild()) {
            if (result.compareTo(root) == 0) {
                root = new Node<T>();
                result = root;
            } else {
                if (result.isLeftChild()) {
                    result = result.getParent();
                    result.setLeftChild(null);
                } else {
                    result = result.getParent();
                    result.setRightChild(null);
                }
            }
        } else if (result.hasOneChild()) {
            if (result.compareTo(root) == 0) {
                if (result.hasLeftChild()) {
                    root = result.getLeftChild();
                } else {
                    root = result.getRightChild();
                }
                root.setParent(null);
                result = root;
            } else {
                if (result.isLeftChild()) {
                    if (result.hasLeftChild()) {
                        result.getParent().setLeftChild(result.getLeftChild());
                    } else {
                        result.getParent().setLeftChild(result.getRightChild());
                    }
                    result = result.getParent();
                    result.getLeftChild().setParent(result);
                } else {
                    if (result.hasLeftChild()) {
                        result.getParent().setRightChild(result.getLeftChild());
                    } else {
                        result.getParent().setRightChild(result.getRightChild());
                    }
                    result = result.getParent();
                    result.getRightChild().setParent(result);
                }
            }
        } else {
            Node<T> successor = result.getRightChild();
            while (successor.hasLeftChild()) {
                successor = successor.getLeftChild();
            }
            successor.setLeftChild(result.getLeftChild());
            result.getLeftChild().setParent(successor);
            if (successor.isLeftChild()) {
                if (successor.hasRightChild()) {
                    successor.getParent().setLeftChild(successor.getRightChild());
                    successor.getRightChild().setParent(successor.getParent());
                } else {
                    successor.getParent().setLeftChild(null);
                }
                successor.setRightChild(result.getRightChild());
                result.getRightChild().setParent(successor);
            }
            if (result == root) {
                successor.setParent(null);
                root = successor;
                result = root;
            } else {
                if (result.isLeftChild()) {
                    successor.setParent(result.getParent());
                    result.getParent().setLeftChild(successor);
                } else {
                    successor.setParent(result.getParent());
                    result.getParent().setRightChild(successor);
                }
                result = successor.getParent();
            }
        }
        --numElements;
        return result;
    }
    
    @Override
    public String toString() {
        String result = "";
        Iterator<T> iter = iteratorLevelOrder();
        while (iter.hasNext()) {
            result += (iter.next() + "\n");
        }
        return result;
    }
    
    protected Iterator<T> iteratorLevelOrder() {
        LinkedList<Node<T>> nodeList = new LinkedList<Node<T>>();
        LinkedList<T> elementList = new LinkedList<T>();
        Node<T> node;
        nodeList.addLast(root);
        while (!nodeList.isEmpty()) {
            node = nodeList.removeFirst();
            if (!node.isEmpty()) {
                elementList.addLast(node.getData());
                if (node.hasLeftChild()) {
                    nodeList.addLast(node.getLeftChild());
                }
                if (node.hasRightChild()) {
                    nodeList.addLast(node.getRightChild());
                }
            }
        }
        return elementList.iterator();
    }
}